package com.piece.core.web.filter;

import com.piece.core.framework.constant.ExceptionConstants;
import com.piece.core.framework.util.collection.ArrayUtil;
import com.piece.core.framework.util.file.IoUtil;
import com.piece.core.framework.util.net.HtmlFilter;
import com.piece.core.framework.util.object.ObjectUtil;
import com.piece.core.framework.util.basic.I18nUtil;
import com.piece.core.framework.util.string.StringUtil;
import com.piece.core.web.util.ServletUtil;
import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

public class XssRequestWrapper extends HttpServletRequestWrapper {

    /**
     * 基于线程级别的 HTMLFilter 对象，因为它线程非安全
     */
    private static final ThreadLocal<HtmlFilter> HTML_FILTER = ThreadLocal.withInitial(() -> {
        HtmlFilter htmlFilter = new HtmlFilter();
        return htmlFilter;
    });

    private List<String> illegalParam = new ArrayList<>();

    public XssRequestWrapper(HttpServletRequest request, String[] illegalParam) {
        super(request);
        this.initIllegalParam(illegalParam);
    }

    public void initIllegalParam(String[] params) {
        if (null != params) {
            this.illegalParam.addAll(Arrays.asList(params));
        }
    }

    private static String filterXss(String content) {
        if (StringUtil.isEmpty(content)) {
            return content;
        }
        return HTML_FILTER.get().filter(content);
    }

    private String filterParam(String content) {
        for (int i = 0; i < this.illegalParam.size(); i++) {
            if (null != content && content.toLowerCase().indexOf(this.illegalParam.get(i)) != -1) {
                throw new IllegalArgumentException(I18nUtil.message(ExceptionConstants.ILLEGAL, content));
            }
        }

        return content;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return filterXss(value);
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        // 如果非 json 请求，不进行 Xss 处理
        if (!ServletUtil.isAjaxRequest(this)) {
            return super.getInputStream();
        }

        // 读取内容，并过滤
        String content = IoUtil.toString(super.getInputStream(), StandardCharsets.UTF_8);
        content = filterXss(content);
        final ByteArrayInputStream newInputStream = new ByteArrayInputStream(content.getBytes());
        // 返回 ServletInputStream
        return new ServletInputStream() {

            @Override
            public int read() {
                return newInputStream.read();
            }

            @Override
            public boolean isFinished() {
                return true;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setReadListener(ReadListener readListener) {}

        };
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        return filterParam(value);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (ArrayUtil.isEmpty(values)) {
            return values;
        }
        // 过滤处理
        for (int i = 0; i < values.length; i++) {
            values[i] = filterParam(values[i]);
        }
        return values;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> valueMap = super.getParameterMap();
        if (ObjectUtil.isEmpty(valueMap)) {
            return valueMap;
        }
        // 过滤处理
        for (Map.Entry<String, String[]> entry : valueMap.entrySet()) {
            String[] values = entry.getValue();
            for (int i = 0; i < values.length; i++) {
                values[i] = filterParam(values[i]);
            }
        }
        return valueMap;
    }
}
